#!/usr/bin/env python2
from __future__ import division

from pylab import *
from numpy import * #unnecessary, just for pydev auto-complete
from matching_pursuit import *
from fJ_generator import *
from scipy import optimize
from scikits.audiolab import wavread
from scikits.audiolab import play
from levenberg_marquardt import ooopti
#from enthought.mayavi.mlab import *
import os

def rotate(l, n):
    return l[n:] + l[:n]

framesize = 44100 / 20

class CLASS:
    def __init__(self):
            t = linspace(0, 1, framesize, endpoint=False)

            self.Bspline = empty((framesize, 4))
            self.Bspline[:, -2] = 0 + 0 * (t ** 1) + 0 * (t ** 2) + 1 * (t ** 3)
            self.Bspline[:, -1] = 1 + 3 * (t ** 1) + 3 * (t ** 2) - 3 * (t ** 3)
            self.Bspline[:, 0] = 4 + 0 * (t ** 1) - 6 * (t ** 2) + 3 * (t ** 3)
            self.Bspline[:, 1] = 1 - 3 * (t ** 1) + 3 * (t ** 2) - 1 * (t ** 3)
            self.Bspline /= 6

            self.Bspline_w = self.Bspline * t.reshape(-1, 1)

            self.t = sympy.Symbol('t')
            self.p = sympy.Symbol('p')
            self.A = [None] * 100
            self.A[-10:] = [sympy.Symbol('A_m' + str(i)[1:]) for i in range(-10, 0)]
            self.A[:10] = [sympy.Symbol('A_' + str(i)) for i in range(0, 10)]
            self.w = [None] * 100
            self.w[-10:] = [sympy.Symbol('w_m' + str(i)[1:]) for i in range(-10, 0)]
            self.w[:10] = [sympy.Symbol('w_' + str(i)) for i in range(0, 10)]
            self.B_A = [None] * 100
            self.B_A[-10:] = [sympy.Symbol('B_A_m' + str(i)[1:]) for i in range(-10, 0)]
            self.B_A[:10] = [sympy.Symbol('B_A_' + str(i)) for i in range(0, 10)]
            self.B_w = [None] * 100
            self.B_w[-10:] = [sympy.Symbol('B_w_m' + str(i)[1:]) for i in range(-10, 0)]
            self.B_w[:10] = [sympy.Symbol('B_w_' + str(i)) for i in range(0, 10)]

            self.fJ_1 = fJ_generator(sympy.log(1 + sympy.exp(self.A[-1] * self.B_A[1] + self.A[0] * self.B_A[0] + self.A[1] * self.B_A[-1] + self.A[2] * self.B_A[-2])) *
                                     sympy.sin(self.p + self.w[-1] * self.B_w[1] + self.w[0] * self.B_w[0] + self.w[1] * self.B_w[-1] + self. w[2] * self.B_w[-2]),

                                     [self.A[-1], self.A[0], self.A[1], self.A[2], self.p, self.w[-1], self.w[0], self.w[1], self.w[2]],

                                       {self.t : linspace(0, 1, framesize, endpoint=False),
                                       self.B_A[-2] : self.Bspline[:, -2],
                                       self.B_A[-1] : self.Bspline[:, -1],
                                       self.B_A[0] : self.Bspline[:, 0],
                                       self.B_A[1] : self.Bspline[:, 1],
                                       self.B_w[-2] : self.Bspline_w[:, -2],
                                       self.B_w[-1] : self.Bspline_w[:, -1],
                                       self.B_w[0] : self.Bspline_w[:, 0],
                                       self.B_w[1] : self.Bspline_w[:, 1]})
    def optimize_1(self, x0, data):
        (result, success, r1, r2, r3) = optimize.leastsq(self.fJ_1.f, x0, args=(data), Dfun=self.fJ_1.J, full_output=True)
        return result
    def optimize_l1(self, l_A, l_w, x0, data):

        self.fJ_l1 = fJ_generator(sympy.log(1 + sympy.exp((l_A[-1] * self.B_A[1] + l_A[0] * self.B_A[0] + l_A[1] * self.B_A[-1] + self.A[2] * self.B_A[-2])*1))/1 *
                                  sympy.sin(self.p + l_w[-1] * self.B_w[1] + l_w[0] * self.B_w[0] + l_w[1] * self.B_w[-1] + self.w[2] * self.B_w[-2]),

                                  [self.A[2], self.p, self.w[2]],

                                  {self.t : linspace(0, 1, framesize, endpoint=False),
                                   self.B_A[-2] : self.Bspline[:, -2],
                                   self.B_A[-1] : self.Bspline[:, -1],
                                   self.B_A[0] : self.Bspline[:, 0],
                                   self.B_A[1] : self.Bspline[:, 1],
                                   self.B_w[-2] : self.Bspline_w[:, -2],
                                   self.B_w[-1] : self.Bspline_w[:, -1],
                                   self.B_w[0] : self.Bspline_w[:, 0],
                                   self.B_w[1] : self.Bspline_w[:, 1]})
        (result, success, r1, r2, r3) = optimize.leastsq(self.fJ_l1.f, x0, args=(data), Dfun=self.fJ_l1.J, full_output=True)
        return result

class DATA:
    def __init__(self, file):
        wav, fs, enc = wavread(file)
        self.data = wav.T[0, :]
    def getframe(self, location):
        return self.data[location:location + framesize]
if __name__ == '__main__':
    _X_ = CLASS()
    _DATA_ = DATA("/home/kkb110/music/g.wav")
    startingpoint = 15.8 * 44100
    data = _DATA_.getframe(startingpoint)

    mp = MatchingPursuit(framesize, 12.5, 120 * 10, 120)
    mp_A, mp_w, mp_p = mp.matching_pursuit(data)
    mp_p += pi / 2
    x0 = [mp_A, mp_A, mp_A, mp_A, mp_p, mp_w, mp_w, mp_w, mp_w]
    result = _X_.optimize_1(x0, data)
    plot(data)
    plot(_X_.fJ_1.f(result))
#    plot(data - _X_.fJ_1.f(result))
    show()

    aaa1 = _X_.fJ_1.f([0, 0, 0, result[0], 9, 0, 0, 0, result[5]], zeros(framesize))
    aaa2 = _X_.fJ_1.f([0, 0, result[0], result[1], 9, 0, 0, result[5], result[6]], zeros(framesize))
    aaa3 = _X_.fJ_1.f([0, result[0], result[1], result[2], 9, 0, result[5], result[6], result[7]], zeros(framesize))
    aaa4 = _X_.fJ_1.f([ result[0], result[1], result[2], result[3], 9, result[5], result[6], result[7], result[8]], zeros(framesize))
    aaa5 = _X_.fJ_1.f([ result[1], result[2], result[3], 0, 9, result[6], result[7], result[8], 0], zeros(framesize))
    aaa6 = _X_.fJ_1.f([ result[2], result[3], 0, 0, 9, result[7], result[8], 0, 0], zeros(framesize))
    aaa7 = _X_.fJ_1.f([ result[3], 0, 0, 0, 9, result[8], 0, 0, 0], zeros(framesize))
    
    
    plot(hstack([aaa4]))
    show()
    toplot = hstack([aaa1, aaa2, aaa3, aaa4, aaa5, aaa6, aaa7])
    plot(toplot)
    print result
    show()
    data2 = _DATA_.getframe(startingpoint + 1 * framesize)
    l_A = rotate(result[1:4], 2)
    l_w = rotate(result[6:9], 2)
#    l_A = result[1:4]
#    l_w = result[6:9]
    result2 = _X_.optimize_l1(l_A, l_w, [1, 0, 0], data2)
#    plot(data2)
#    plot(_X_.fJ_l1.f([0, 0, 0]))
##    plot(_X_.fJ_l1.f(result2) - data2)
#    show()























